###############################################################################
#                                                                             #
#    L3q - Light, light, lightweight queue                                    #
#    Copyright (C) 2023-2024  Marcus Pedersén marcus@marcux.org               #
#                                                                             #
#    This program is free software: you can redistribute it and/or modify     #
#    it under the terms of the GNU General Public License as published by     #
#    the Free Software Foundation, either version 3 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
#    This program is distributed in the hope that it will be useful,          #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
#    GNU General Public License for more details.                             #
#                                                                             #
#    You should have received a copy of the GNU General Public License        #
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.    #
#                                                                             #
###############################################################################


'''
This module contains functionality
for the L3QWorker.

'''



import signal
import sys
import os
import shutil
from libl3q.daemon import sql, queue
from libl3q.daemon.daemon import NodeState, JobState
from subprocess import Popen
from math import ceil
import datetime
import psutil
import threading
import requests
import json
import urllib3
from typing import Any

from multiprocessing import Process, Pipe

class L3QWorker:
    ''''
    The L3QWorker is the process that checks 
    status of nodes and gathers info from them.
    It also updates the queue, starts processes
    in the queue and cancel jobs.
    '''

    def __init__(self, conf):
        '''
        Init the L3QWorker and set up
        required config for process.
        '''

        self.logger = conf.logger
        self.logger.info('Starting L3QWorker ...')
        self.db_conn = sql.init_db(conf)
        self.db_cur = self.db_conn.cursor()
        self.poll_intervall = conf.poll_intervall
        self.user_group = conf.user_group
        self.queue = conf.queue
        self.timer = None
        self.running = 'Done'

        # Disable warnings when using self signed cert
        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        
        signal.signal(signal.SIGTERM, self.terminate)
        
    def run(self):
        '''
        Starts and run required functionality,
        checking node status, updates queue
        an cancel jobs.
        '''

        self.timer = threading.Event()
        self.running = 'Done'
        
        while not self.timer.is_set():
            if(not self.timer.is_set()):
                self.running = 'Update nodes'
                self.__update_node_states()
                self.running = 'Done'

            if(not self.timer.is_set()):
                self.running = 'Update queue'
                self.__update_depend_jobs()
                self.__cancel_jobs()
                # self.__update_running_jobs()
                self.__update_queue()
                self.running = 'Done'

            self.timer.wait(self.poll_intervall)
            if(not self.timer.is_set()):
                self.timer = threading.Event()


    def __update_node_states(self):
        '''
        Connects to all nodes to get current status
        and usage of node. Info is updated in database.
        '''

        get_status = { 'cmd': 'Status', 'validate': '' }
        try:
            self.db_cur.execute('''SELECT validate_host.id, ip, port, hostname, key, l3qd_key, state
            FROM validate_host
            INNER JOIN node
            ON validate_host.id = node.validate
            WHERE host_type = 'node'
            AND
            (state = ? OR state = ? OR state = ? OR state = ?)''',
                                (NodeState.REGISTER, NodeState.ONLINE, NodeState.OFFLINE, NodeState.SOFT_ONLINE))

            for n in self.db_cur:
                try:
                    get_status['validate'] = n['key']
                    json_data = json.dumps(get_status)
                except Exception as e:
                    self.logger.error('Worker: Error encoding data into json string:\n{}'.format(e))
                    continue

                try:
                    resp = requests.post('https://' + n['hostname'] + ':' + str(n['port']) + '/l3qd/request',
                                         data=json_data, verify=False)
                    resp.raise_for_status()
                except:
                    try: 
                        resp = requests.post('https://' + n['ip'] + ':' + str(n['port']) + '/l3qd/request',
                                             data=json_data, verify=False)
                        resp.raise_for_status()
                    except:
                        self.__set_node_offline(n['id'])
                        continue

                try:
                    data = json.loads(resp.text)
                except Exception as e:
                    self.logger.error('Worker: Failed to parse response from node (validate id): {}'.format(n['id']))
                    self.__set_node_offline(n['id'])
                    continue

                if(('NodeStatus' in data['data']) and ('validate' in data)):
                    if(data['validate'] == n['l3qd_key']):
                        self.db_cur.execute('''UPDATE node SET 
                        total_cpus = ?,
                        used_cpus = ? ,
                        total_memory = ?,
                        used_memory = ?,
                        total_swap = ?,
                        used_swap = ?,
                        state = ?
                        WHERE
                        validate = ?''',
                                            (data['data']['NodeStatus']['total_cpus'],
                                             data['data']['NodeStatus']['used_cpus'],
                                             data['data']['NodeStatus']['total_memory'],
                                             data['data']['NodeStatus']['used_memory'],
                                             data['data']['NodeStatus']['total_swap'],
                                             data['data']['NodeStatus']['used_swap'],
                                             NodeState.ONLINE,
                                             n['id']))
                    else:
                        self.logger.error('Worker: Validation failure in answer from node (validate id): {}'.format(n['id']))
                        self.__set_node_offline(n['id'])
                else:
                    self.logger.error('Worker: Data is missing in answer from node (validate id): {}'.format(n['id']))
                    self.__set_node_offline(n['id'])
        except Exception as e:
            self.logger.error('Worker: Failed to get all nodes from database: {}'.format(e))
            

        # Update nodes that have state Maintenance draining
        try:
            self.db_cur.execute('''UPDATE node SET state = ? WHERE alloc_cpus = 0 AND state = ?''',
                                (NodeState.MAINTENANCE, NodeState.MAINTENANCE_DRAIN))
        except Exception as e:
            self.logger.error('Worker: Failed to update states for node with state: {}\n{}'.format(NodeState.MAINTENANCE_DRAIN, e))


    def terminate(self, signum, frame):
        self.logger.info('Worker: Closing down ...')

        if self.timer:
            self.timer.set()

        if(self.running != 'Done'):
            self.logger.info('Worker: Waiting for: {} to finish ...'.format(self.running))

        try:
            self.db_conn.flush()
            self.db_conn.close()
            self.logger.info('Worker: Terminating')
        except:
            pass
            


    def __set_node_offline(self, validateid):
        ''' Sets node offline in database'''

        try:
            cur = self.db_conn.execute('''UPDATE node SET 
used_cpus = NULL,
used_memory = NULL,
used_swap = NULL,
state = ?
WHERE validate = ?''', (NodeState.OFFLINE, validateid))
            if(cur.rowcount == 0):
                self.logger.error('Worker: Failed to update node to state: OFFLINE for node (validate id): {}, zero rows updated.'.format(validateid))
        except Exception as e:
            self.logger.error('Worker: Failed to update node to state: OFFLINE for node (validate id): {}, {}'.format(validateid, e))




    def __update_depend_jobs(self):
        '''
        Checks if depending jobs has exited
        in the correct way and then this job
        will be queued.
        If failure of depending jobs this job
        will also be set as fail.
        '''

        try:
            self.db_cur.execute('''SELECT id, depend FROM job WHERE state = ? AND depend IS NOT NULL AND NOT depend = '' ''',
                                (JobState.DEPEND,))
            deps = self.db_cur.fetchall()

            for d in deps:
                dep_state = ''
                dep_jobs = d['depend'].strip().split(',')
                dep_sql = f'SELECT id, state FROM job WHERE id = {dep_jobs[0]}'
                
                for j in dep_jobs[1:]:
                    dep_sql += f' OR id = {j}'
                    
                self.db_cur.execute(dep_sql)

                no_term = 0
                for j in self.db_cur:
                    stat = j['state']
                    if(stat == JobState.CANCEL or
                       stat == JobState.CANCELED or
                       stat == JobState.ERROR or
                       stat == JobState.NODE_ERROR):
                        dep_state = JobState.ERROR
                        break
                    elif(stat == JobState.TERMINATED):
                        no_term += 1

                if(no_term == len(dep_jobs)):
                    dep_state = JobState.QUEUED

                if(len(dep_state) > 0):
                    self.db_cur.execute('''UPDATE job SET state = ? WHERE id = ?''',
                                        (dep_state , d['id']))
                    
        except Exception as e:
            self.logger.error(f'Worker: Failed to update dependencies job in database: {e}')



    def __cancel_jobs(self):
        '''
        Cancel jobs that client has asked
        to cancel.
        Terminates running process and 
        updates job in database.
        '''

        try:
            self.db_cur.execute('''SELECT id, no_tasks_running, nodes, cores_alloc
            FROM job WHERE state = ?''',
                                (JobState.CANCEL,))

            cancel_jobs = self.db_cur.fetchall()

            for c in cancel_jobs:
                if(c['no_tasks_running'] == 0):
                    self.db_cur.execute('''UPDATE job set state = ? WHERE id = ?''',
                                        (JobState.CANCELED, c['id']))
                else:
                    no_nodes_cancel = 0
                    nodes = c['nodes'].strip().split(',')

                    for n in nodes:
                        self.db_cur.execute('''SELECT hostname, ip, port, key, l3qd_key
                        FROM node
                        JOIN validate_host
                        ON node.validate = validate_host.id
                        WHERE node.id = ?''',
                                            (n,))

                        node = self.db_cur.fetchone()

                        if(node == None):
                            self.logger.error('Worker: Failed to fetch node with nodeid: {} from database.'
                                              .format(n))
                            continue
                        
                        cmd_cancel = { 'cmd': { 'CancelL3qdJob': { 'jobid': int(c['id']) } }, 'validate': node['key'] }

                        try:
                            json_data = json.dumps(cmd_cancel)
                        except Exception as e:
                            self.logger.error('Worker: Error encoding Cancel Cmd data into json string:\n{}'.format(e))
                            continue
                            
                        try:
                            resp = requests.post('https://' + node['hostname'] + ':' + str(node['port']) + '/l3qd/request',
                                                 data=json_data, verify=False)
                            resp.raise_for_status()
                        except:
                            try: 
                                resp = requests.post('https://' + node['ip'] + ':' + str(node['port']) + '/l3qd/request',
                                                     data=json_data, verify=False)
                                resp.raise_for_status()
                            except Exception as e:
                                self.logger.error('Worker: Failed to cancel jobid: {} on nodeid: {}\n{}'
                                                  .format(c['id'], n, e))
                                continue
                        try:
                            node_resp = json.loads(resp)
                        except Exception as e:
                            self.logger.error('Worker: Failed to parse cancel response jobid: {} from nodeid: {}\n{}'
                                              .format(c['id'], n, e))
                            continue

                        if(node_resp['validate'] != node['l3qd_key']):
                            self.logger.error('Worker: Failed to validate key in response fron node: {}, {}, {} '
                                              .format(n, node['hostname'], node['ip']))
                            continue
                                              

                        if('L3qdNodeResponse' in node_resp):
                            if('CancelledTasks' in node_resp['L3qdNodeResponse']['data']):
                                no_nodes_cancel += 1

                                try:
                                    node_alloc = 0
                                    self.db_cur.execute('''SELECT alloc_cpus from node WHERE id = ?''', (n,))
                                    node = self.db_cur.fetchone()

                                    if(node == None):
                                        self.logger.error('Worker: Failed to fetch alloc_cpus from nodeid: {} in database.'
                                                          .format(n))
                                        continue
                                    
                                    if(c['cores_alloc'] >= node['alloc_cpus']):
                                        node_alloc = 0
                                    else:
                                        node_alloc = node['alloc_cpus'] - c['cores_alloc']
                                        
                                    self.db_cur.execute('''UPDATE node SET alloc_cpus = ? WHERE id = ?''', (node_alloc, n))
                                    
                                except Exception as e:
                                    self.logger.error('Worker: Failed to update nodeid: {} allocation for jobid: {} after all tasks has been canceled on node.\n{}'
                                                      .format(n, c['id'], e))

                            else:
                                self.logger.error('Worker: Failed to cancel jobid: {} on nodeid: {}\nNode error: {}'
                                                  .format(c['id'], n, node_resp['L3qdNodeResponse']['data']))
                                break

                    if(no_nodes_cancel == len(nodes)):
                        try: 
                            self.db_cur.execute('''UPDATE job set state = ? WHERE id = ?''',
                                            (JobState.CANCELED, c['id']))
                        except Exception as e:
                            self.logger.error('Worker: Failed to update jobid: {} status after all tasks has been canceled\n{}',
                                              (c['id'], e))

        except Exception as e:
            self.logger.error(f'Worker: Failed to cancel jobs marked for canceling.\n{e}')




    def __update_queue(self) -> None:
        '''
        Update configured queue.
        Sends all queued and depending
        jobs and node info to queue 
        object that returns jobs to
        be launched on the nodes.
        Send jobs to the nodes and
        updates the database.
        '''

        q = eval(f"queue.{self.queue}()")
        job_queue = []
        try:
            self.db_cur.execute('SELECT id,type,init_date,hosts_alloc,cores_alloc,depend,user,state FROM job WHERE state = ?',
                                (JobState.QUEUED,))
            db_job_queue = self.db_cur.fetchall()

            for j in db_job_queue:
                no_tasks = 0
                
                try:
                    self.db_cur.execute('SELECT count(*) as count FROM task WHERE jobid = ?',
                                        (j['id'],))
                    task_count = self.db_cur.fetchone()
                    no_tasks = task_count['count']
                except Exception as e:
                    self.logger.error('Worker: Failed to get all tasks for jobid: {} from database:\n{}'
                                      .format(j['id'],e))
                    
                qj = queue.QJob.from_dict(j, no_tasks)

                if(qj != None):
                    job_queue.append(qj)

        except Exception as e:
            self.logger.error(f'Worker: Failed to get all queued jobs from database: {e}')
            return 
                        
        online_nodes = []
        try:
            self.db_cur.execute('SELECT id,total_cpus,used_cpus,alloc_cpus,total_memory,used_memory,total_swap,used_swap FROM node WHERE state = ?', (NodeState.ONLINE,))

            nodes = self.db_cur.fetchall()

            for n in nodes:
                node = queue.QNode.from_dict(n)

                if(node != None):
                    online_nodes.append(node)

        except Exception as e:
            self.logger.error(f'Worker: Failed to get all online nodes from database: {e}')

        run = q.run_queue(job_queue, online_nodes)
        for r in run:
            try:
                self.db_cur.execute('SELECT * FROM task WHERE jobid = ? AND start_time IS NULL', (r.job.jobid,))
                tasks = self.db_cur.fetchall()
                
            except Exception as e:
                self.logger.error(f'Worker: Failed to get tasks for jobid: {r.job.jobid} from database: {e}')



            used_tasks = []
            for n in r.node_ids:
                node_tasks = []

                i = 1
                while(i <= r.job.cores_alloc):

                    for t in tasks:
                        if(t['id'] not in used_tasks):
                            task = {
                                'taskid': t['id'],
                                'command': t['command'],
                                'working_directory': t['workdir']
                            }

                            node_tasks.append(task)
                            used_tasks.append(t['id'])
                            break
                    
                    i += 1
                                     

                node_job = {
                    'cmd':
                    {
                        'ExecuteL3qdJob':
                        {
                            'job':
                            { 
                                'jobid': r.job.jobid,
                                'cores_alloc': r.job.cores_alloc,
                                'memory_alloc': 0,
                                'exec_type': r.job.job_type,
                                'tasks': node_tasks,
                                'user': r.job.user,
                                'group': '',
                            }
                        }
                    },
                    'validate': ''
                }

                if(self.__send_job_to_node(n, node_job)):
                    try:
                        self.db_cur.execute('''SELECT total_cpus,alloc_cpus
                        FROM node
                        WHERE id = ?''', (n,))

                        node = self.db_cur.fetchone()
                        if(node['alloc_cpus'] == None):
                            tot_alloc = node_job['cmd']['ExecuteL3qdJob']['job']['cores_alloc']
                        else:
                            tot_alloc = node['alloc_cpus'] + node_job['cmd']['ExecuteL3qdJob']['job']['cores_alloc']
                            
                        if tot_alloc > node['total_cpus']:
                            tot_alloc = node['total_cpus']

                        self.db_cur.execute('''UPDATE node SET alloc_cpus = ? WHERE id = ?''',
                                            (tot_alloc, n))
                        
                    except Exception as e:
                        self.logger.error('Worker: Failed to update database with alloc_cpus on node id: {}\n{}'.format(n,e))

                    try:
                        self.db_cur.execute('''SELECT no_tasks_running
                        FROM job
                        WHERE id = ?''', (node_job['cmd']['ExecuteL3qdJob']['job']['jobid'],))

                        job = self.db_cur.fetchone()
                        running = job['no_tasks_running'] + len(node_job['cmd']['ExecuteL3qdJob']['job']['tasks'])

                        self.db_cur.execute('''UPDATE job
                        SET
                        start_date = ?,
                        no_tasks_running = ?,
                        state = ?
                        WHERE id = ?''',
                                            (datetime.datetime.now().isoformat(timespec='seconds', sep=' '),
                                             running, JobState.RUNNING, node_job['cmd']['ExecuteL3qdJob']['job']['jobid']))
                        
                    except Exception as e:
                        self.logger.error('Worker: Failed to update database on job id: {} after sent to node.\n{}'
                                          .format(node_job['cmd']['ExecuteL3qdJob']['job']['jobid'],e))

                    for t in node_job['cmd']['ExecuteL3qdJob']['job']['tasks']:
                        try:
                            self.db_cur.execute('''UPDATE task SET nodeid = ? WHERE id = ?''',
                                                (n, t['taskid']))
                        except Exception as e:
                            self.logger.error('Worker: Failed to update database with node id: {} on task id: {}\n{}'
                                              .format(n, t['taskid'], e))





    def __send_job_to_node(self, node_id: int, job: dict[str, Any]) -> bool:
        '''
        Sends a new job to specified node.
        Tries to connect to node both with
        IP number and hostname. If both fails
        then job sending failed.
        Returns true if job is successfully
        sent to node.
        '''
        
        try:
            self.db_cur.execute('''SELECT validate_host.id, ip, port, hostname, key, l3qd_key
            FROM validate_host
            INNER JOIN node
            ON validate_host.id = node.validate
            WHERE node.id = ?''', (node_id,))

            node = self.db_cur.fetchone()
        except Exception as e:
            self.logger.error('Worker: Failed to get node with id: {} from database:\n{}'.format(node_id,e))
            return False

        if(node == None):
            return False

        job['validate'] = node['key']
        try:
            json_data = json.dumps(job)
        except Exception as e:
            self.logger.error('Worker: Error encoding data into json string:\n{}'.format(e))
            return False
                
        try:
            resp = requests.post('https://' + node['hostname'] + ':' + str(node['port']) + '/l3qd/request',
                                 data=json_data, verify=False)
            resp.raise_for_status()
        except:
            try: 
                resp = requests.post('https://' + node['ip'] + ':' + str(node['port']) + '/l3qd/request',
                                     data=json_data, verify=False)
                resp.raise_for_status()
            except Exception as e:
                self.logger.error('Worker: Failed to send jobid: {} to node id: {}, {}'
                                  .format(job['cmd']['ExecuteL3qdJob']['job']['jobid'],
                                          node_id, e))
                return False

        try:
            data = json.loads(resp.text)
        except Exception as e:
            self.logger.error('Worker: Failed to parse response from node: {} when sending job with id: {}'
                              .format(node_id, job['cmd']['ExecuteL3qdJob']['job']['jobid']))
            return False
        if(not data):
            self.logger.error('Worker: Error response from node: {} when sending job with id: {}\nEmpty respone from node.'
                              .format(node_id, job['cmd']['ExecuteL3qdJob']['job']['jobid']))
            return False
        elif('Error' in data['data']):
            self.logger.error('Worker: Error response from node: {} when sending job with id: {}\nNode error: {}'
                              .format(node_id, job['cmd']['ExecuteL3qdJob']['job']['jobid'], data['data']['Error']))
            return False
        elif(('TaskExecuted' in data['data']) and ('validate' in data)):
            if(data['validate'] == node['l3qd_key']):
                if(data['data']['TaskExecuted']['jobid'] == job['cmd']['ExecuteL3qdJob']['job']['jobid']):
                    sent_tasks = []
                    for t in  job['cmd']['ExecuteL3qdJob']['job']['tasks']:
                        sent_tasks.append(t['taskid'])
                        
                    resp_tasks = data['data']['TaskExecuted']['task']

                    if(sorted(sent_tasks) == sorted(resp_tasks)):
                        return True
                    else:
                        self.logger.error('Worker: Wrong taskids: {} in answer from node: {} when sending job with id: {} and taskids: {}'
                              .format(resp_tasks,
                                      node_id,
                                      job['cmd']['ExecuteL3qdJob']['job']['jobid'],
                                      sent_tasks))
                        return False
                else:
                    self.logger.error('Worker: Wrong jobid: {} in answer from node: {} when sending job with id: {}'
                              .format(data['data']['TaskExecuted']['jobid'],
                                      node_id,
                                      job['cmd']['ExecuteL3qdJob']['job']['jobid']))
                    return False
            else:
                self.logger.error('Worker: Validation failure from node: {} when sending job with id: {}'
                              .format(node_id, job['cmd']['ExecuteL3qdJob']['job']['jobid']))
                return False
                

        return True
